import pygame 
import random

SCREEN_SIZE = [1000, 800]
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
CELL_RADIUS = 20
SNAKE_COLOR = (0, 0, 0)
UPDATE = pygame.USEREVENT + 1
FOOD = pygame.USEREVENT + 2
FOOD_COLOR = (0, 100, 0)
FOOD_RADIUS = 15
LIGHT_GREY = (100,100,100)
MSG_POSITION = (300, 350)

# 游戏初始化
def game_init():
    pygame.init()
    pygame.display.set_caption("Snake 贪吃蛇")
    pygame.time.set_timer(UPDATE, 300) # 每500毫秒执行一遍视图更新
    pygame.time.set_timer(FOOD, 2000) # 每500毫秒执行一遍食物更新

# 游戏类
class Game:
    def __init__(self) -> None:
        self.running = True 
        self.screen = pygame.display.set_mode(SCREEN_SIZE)
        self.snake = Snake()
        self.food = None
        self.message = None
        self.images = {}


# 蛇的每一个点的位置
class Cell():
    def __init__(self, x, y) -> None:
        self.x = x
        self.y = y

    def to_tuple(self):
        return self.x, self.y

    def copy(self):
        return Cell(self.x, self.y)
    
    def update(self, direction):
        if direction == "U":
            self.y -= CELL_RADIUS *2
        elif direction == "D":
            self.y += CELL_RADIUS *2
        elif direction == "L":
            self.x -= CELL_RADIUS *2
        else:
            self.x += CELL_RADIUS *2

# 蛇类
class Snake():
    def __init__(self) -> None:
        cell_diameter = CELL_RADIUS * 2
        x = (SCREEN_SIZE[0]/cell_diameter//2) * cell_diameter
        y = (SCREEN_SIZE[1]/cell_diameter//2) * cell_diameter

        self.cell_size = CELL_RADIUS 
        self.color = SNAKE_COLOR
        self.direction = "L"
        self.body = [Cell(x, y)]

    # 按键按下方向之后更新蛇的身体
    def update(self):
        head = self.body[0].copy() # 1. 复制头
        self.body.pop() # 2. 去掉尾
        head.update(self.direction) # 3. 计算按下按键后的下一个位置
        self.body.insert(0, head) # 插入body

# 更新
def update():
    check_key_pressed() # 检查按键的方向
    check_food() # 检测蛇的位置和食物的位置
    check_head_body_collision() # 检测头和身体的碰撞
    check_out_boundary() # 检测蛇有没有出边界
    check_win() # 检测游戏是否赢了
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            game.running = False # 终止游戏
        if event.type == UPDATE:
            game.snake.update() # 根据direction更新蛇的body
        if event.type == FOOD:
            generate_food()
            
# 绘画 
def draw():
    game.screen.fill(WHITE)
    if game.running:
        draw_snake() # 画蛇的body
        draw_food() # 画出蛇的位置
    else:
        draw_restart() # 重新开始游戏
    pygame.display.flip()

# 画蛇的身体
def draw_snake():
    for cell in game.snake.body:
        pygame.draw.circle(game.screen, game.snake.color, cell.to_tuple(), game.snake.cell_size)

# 画出食物的位置
def draw_food():
    if game.food is not None:
        pygame.draw.circle(game.screen, FOOD_COLOR, game.food, FOOD_RADIUS)

# 重新开始游戏
def draw_restart():
    font = pygame.font.Font(None, 30)
    restart_msg = font.render(game.message, True, LIGHT_GREY)
    game.screen.blit(restart_msg, MSG_POSITION)
    # game.screen.blit(restart_msg, (10,10))

# 随机生成食物位置
def generate_food():
    while game.food is None:
        cell_diameter = CELL_RADIUS * 2
        rand_x = random.randint(1, (SCREEN_SIZE[0]/cell_diameter)-1)* cell_diameter
        rand_y = random.randint(1, (SCREEN_SIZE[1]/cell_diameter)-1)* cell_diameter
        if (rand_x, rand_y) not in [cell.to_tuple() for cell in game.snake.body]:
            game.food = (rand_x, rand_y)

# 检测按下的是什么键
def check_key_pressed():
    pressed_key = pygame.key.get_pressed()
    if pressed_key[pygame.K_UP] and not game.snake.direction == "D":
        game.snake.direction = "U"
    if pressed_key[pygame.K_DOWN] and not game.snake.direction == "U":
        game.snake.direction = "D"
    if pressed_key[pygame.K_LEFT] and not game.snake.direction == "R":
        game.snake.direction = "L"
    if pressed_key[pygame.K_RIGHT] and not game.snake.direction == "L":
        game.snake.direction = "R"

# 检测蛇的位置和食物的位置
def check_food():
    if game.food is not None and is_snake_food_collide(): # 检测食物位置和蛇的位置
        cell = Cell(game.food[0], game.food[1])
        game.snake.body.insert(0, cell)
        game.food = None

# 检测蛇头和食物是否在同一个位置(吃)
def is_snake_food_collide():
    head = game.snake.body[0].copy()
    head.update(game.snake.direction)
    return (game.food[0] - head.x) == 0 and (game.food[1] - head.y) == 0

# 检测头部有没有碰撞到身体
def check_head_body_collision():
    head = game.snake.body[0]
    body = game.snake.body
    try:
        index = [cell.to_tuple() for cell in body[1:]].index(head.to_tuple())
    except ValueError:
        index = -1
    if len(game.snake.body) > 1 and index > -1:
        game.snake.body = game.snake.body[:index]

# 检测蛇是否出边界
def check_out_boundary():
    cell_diameter = CELL_RADIUS *2
    out_boundary_x_cells = [cell for cell in game.snake.body 
                                    if cell.x < cell_diameter or cell.x > SCREEN_SIZE[0]-cell_diameter]
    out_boundary_y_cells = [cell for cell in game.snake.body 
                                    if cell.y < cell_diameter or cell.y > SCREEN_SIZE[1]-cell_diameter]
    if len(out_boundary_x_cells) > 0 or len(out_boundary_y_cells) > 0:
        game.running = False
        game.message = "Snake is out of boundary. Should restart? Y / N"

# 检测游戏是否赢了
def check_win():
    length = len(game.snake.body)
    cell_diameter = CELL_RADIUS * 2
    max_cell_x = SCREEN_SIZE[0] // cell_diameter
    max_cell_y = SCREEN_SIZE[1] // cell_diameter
    if length == max_cell_x * max_cell_y:
        game.running = False
        game.message = "You WIN the game. Should restart? Y / N "

# 检测是否要重新开始游戏
def check_restart():
    while not game.running:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                return
            if event.type == pygame.KEYUP and event.key == pygame.K_y:
                game.snake = Snake()
                game.running = True
                    
            if event.type == pygame.KEYUP and event.key == pygame.K_n:
                return

if __name__=="__main__":
    game_init()
    game = Game()
    while game.running:
        update()
        draw()
        check_restart()
                
    pygame.quit()

# def load_images():
#     head_img = pygame.image.load("./dasdf")
#     head_img = pygame.transform.scale(head_img, (radius, radius))
#     game.images = {
#         "head_img": head_img,
#         "food":food_img,
#     }

# game.screen.blit(game.images["head_img"], game.snake.body[0].to_tuple())







